import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import warnings

warnings.filterwarnings('ignore')
import tensorflow as tf

tf.compat.v1.logging.set_verbosity(40)

from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout, BatchNormalization, Activation, AveragePooling2D, GlobalAveragePooling2D

(x_train,y_train),(x_test,y_test)=tf.keras.datasets.cifar10.load_data()
x_train=x_train.reshape([-1,32,32,3])/255
x_test=x_test.reshape([-1,32,32,3])/255

def conv2dCell(input, filters, num_row, num_col, strides=(1, 1)):
    x = Conv2D(filters, (num_row, num_col), strides=strides, padding='same')(input)
    x = BatchNormalization()(x)
    x = Activation('relu')(x)
    return x

def InceptionBlock(input, filters, strides=(1, 1)):
    branch_1 = conv2dCell(input, filters, 1, 1, strides=strides)
    branch_2 = conv2dCell(input, filters, 1, 1, strides=strides)
    branch_2 = conv2dCell(branch_2, filters, 3, 3)
    branch_3 = conv2dCell(input, filters, 1, 1, strides=strides)
    branch_3 = conv2dCell(branch_3, filters, 5, 5)
    branch_4 = AveragePooling2D((3, 3), strides=(1, 1), padding='same')(input)
    branch_4 = conv2dCell(branch_4, filters, 1, 1, strides=strides)
    output = tf.keras.layers.concatenate(
        [branch_1, branch_2, branch_3, branch_4], axis=3)
    return output

inputs = tf.keras.Input([32, 32, 3])
x = conv2dCell(inputs, 16, 3, 3, strides=(1, 1))  # out:（32，32，16）

num_blocks = 2
init_ch = 16
out_ch = init_ch
for _ in range(num_blocks):
    out_ch *= 2
    for id in range(2):
        if id == 0:
            x = InceptionBlock(x, out_ch, strides=(2, 2))
        else:
            x = InceptionBlock(x, out_ch, strides=(1, 1))

x = GlobalAveragePooling2D()(x)
outputs = Dense(10, activation='softmax')(x)

model = tf.keras.Model(inputs, outputs)

model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(),
              optimizer=tf.keras.optimizers.Adam(lr=0.001),
              metrics=['accuracy'])

model.summary()

history = model.fit(x_train, y_train, batch_size=64, epochs=1, validation_split=0.3)

